"""
This module provides classes and functions for an individual based modeling 
approach to the question of whether an MK-Wolbachia infected population is 
selected for shifts in individual sex ratio (assumed to be under the control 
of females).

'ibmrp' is an acronym for 'individual based modeling of reproductive parasites'
"""
__version__ = 0.04
__author__ = 'M. Flor'

import numpy as np
import numpy.random as npr
import matplotlib as mpl
import matplotlib.pyplot as plt
import copy, string, itertools, bisect, math, time, pprint, textwrap

def build_allele_names(n):
    return ["a{}".format(i) for i in range(n)]
    
def build_phenotype_map(alleles):
    return dict( zip(alleles,range(len(alleles))) )

def init_parameters(**kwargs):
    """Initialize parameters."""
    # first, fill parameter dictionary with default values:
    pars = {'female_fecundity': 10, \
            'mk_lethality': 1., \
            'mk_transmission_rate': 0.95, \
            'reallocation_efficiency': 0.2, \
            'mutation_rate': 1e-5, \
            'probability_of_large_effect_mutation': 0.001, \
            'population_capacity': 100, \
            'base_fitness': 1.}
    # update default parameters with provided arguments:
    pars.update(**kwargs)
    # calculate the remaining parameters (these should not be provided explicitely!):
    pars['number_of_alleles'] = pars['female_fecundity'] + 1
    pars['alleles'] = build_allele_names(pars['number_of_alleles'])
    pars['phenotype_map'] = build_phenotype_map(pars['alleles'])
    return pars

# default module parameters:
_PARAMETERS = init_parameters()

def repr_parameters(pars):
    s = ''
    parnames = ['female_fecundity', 'mk_lethality', 'mk_transmission_rate', 'reallocation_efficiency', \
                'mutation_rate', 'probability_of_large_effect_mutation', 'population_capacity', \
                'base_fitness']
    for name in parnames:
        s += "{}: {}\n".format(name, pars[name])
    return s

def print_parameters(pars):
    print repr_parameters(pars)

def save_parameters(pars, filename='parameters.txt'):
    """Save parameters to file in a human readable format."""
    parfile = open(filename, 'w')
    parfile.write(repr_parameters(pars))
    parfile.close()
    
def load_parameters(filename):
    """Load parameters from file."""
    parfile = open(filename, 'r')
    pars = {}
    lines = parfile.readlines()
    for line in lines:
        p,v = line.split(':')
        pars[p] = eval(v)
    parfile.close()
    return initParameters(**pars)
     
def random_choice(seq):
    return seq[npr.randint(len(seq))]
    
def random_genotype(alleles):
    return [random_choice(alleles), random_choice(alleles)]

def shift_mutate(a, alleles):
    """Shift the mutated allele to a neighboring allele."""
    i = alleles.index(a)
    if npr.random() < 0.5:                # shift left
        if not i == 0:
            i -= 1
    else:                               # shift right
        if not i == len(alleles)-1:
            i += 1
    return alleles[i]

def jump_mutate(a, alleles):
    """Any allele different than the one to mutate is possible."""
    mutant_alleles = alleles[:]
    mutant_alleles.remove(a)
    return random_choice( mutant_alleles )
        
def mutate(a, alleles, mode='mixed', p_jump=0.001):
    if mode=='jump':
        return jump_mutate(a, alleles)
    elif mode=='shift':
        return shift_mutate(a, alleles)
    elif mode=='mixed':
        if npr.random() < p_jump:    # probability that mutation is of large effect
            return jump_mutate(a, alleles)
        else:
            return shift_mutate(a, alleles)
    else:
        raise ValueError, "mutation mode unknown"

def fitness_compensation(x, F, b, rmin=None):
    """
    Fitness compensation / resource reallocation based 
    
    x    - number of killed sons (integer value!)
    F    - total clutch size (female fecundity)
    rmin - resource threshold to viability
    b    - reallocation efficiency
    
    The fitness increase due to resource 
    Example:
    >>> from numpy import arange
    >>> from ibm_v4 import fitness_compensation
    >>> from scitools.aplotter import plot as aplot
    >>> F = 100
    >>> b = 0.5
    >>> rmin = 0.001    # default value
    >>> x = arange(0,F+1)
    >>> y = fitness_compensation(x, F, b, rmin)
    >>> aplot(x, y, dot='o', min_y=0, plot_slope=False)    # this plots fitness compensation per surviving offspring vs number of killed sons
       |                                                                            
       +10                                                                      o   
       |                                                                            
       |                                                                       o    
       |                                                                            
       |                                                                      o     
       |                                                                            
       |                                                                      o     
       |                                                                    oo      
       |                                                                    o       
       |                                                                  oo        
       |                                                                 o          
       |                                                              ooo           
       |                                                         ooooo              
       |                                                  oooooooo                  
       |                                 ooooooooooooooooo                          
       oooooooooooooooooooooooooooooooooo                                           
       |                                                                            
    ---+0-----------------------------------------------------------------------+---
       +0                                                                   +100    
    """
    if not rmin:
        rmin = 1./(10*F)            # default value: optimal resource level is 10 times the threshold
    a = (1-F*rmin)*(F-x)
    numerator   = a + b*x
    denominator = a + b*x*F*rmin
    return numerator / denominator

def fitness_compensation_old(x, b, t):
    """\
    x - sex ratio  (float value!)
    b - reallocation efficiency
    t - transmission rate of MK
    """
    return b*x*t/(1-x*t)

def treeindent(s, pos):
    lines = s.splitlines()
    prefix = '                |'+' '*(pos-5)
    inline = lines[0]+'\n'
    for l in lines[1:]:
        inline += "{}{}\n".format(prefix,l)
    return inline.rstrip()

def weighted_random_generator(weights, replacement=False):
    """Factory function to create a weighted random generator."""
    if replacement:
        return WeightedRandomGeneratorWithReplacement(weights)
    else:
        return WeightedRandomGeneratorWithoutReplacement(weights)

def random_individual(parameters=_PARAMETERS, mk=None):
    alleles = parameters['alleles']
    if npr.random() < 0.5:
        infection = None
    else:
        if mk == None:
            infection = Malekiller(parameters=parameters)
        else:
            infection = mk
    if npr.random() < 0.5:
        return Male(random_genotype(alleles), infection, parameters=parameters)
    else:
        return Female(random_genotype(alleles), infection, parameters=parameters)
        
def random_population(N, parameters=_PARAMETERS):
    males = []
    females = []
    mk = Malekiller(parameters=parameters)
    for i in range(N):
        individual = random_individual(parameters, mk)
        if individual.sex == 'male':
            males.append(individual)
        else:
            females.append(individual)
    return Population(males, females, parameters=parameters)

#def uniform_population(gnum=1, mk=Malekiller()):
#    """gnum - number of males/females per genotype"""
#    alleles = _ALLELES
#    nA = _NUMBER_OF_ALLELES
#    females = []
#    males = []
#    for i in range(nA):
#        for j in range(nA):
#            for k in range(gnum):
#                if npr.random() < 0.5:
#                    infection = None
#                else:
#                    infection = mk
#                male = Male([alleles[i],alleles[j]],infection)
#                if npr.random() < 0.5:
#                    infection = None
#                else:
#                    infection = mk
#                female = Female([alleles[i],alleles[j]],infection)
#                males.append(male)
#                females.append(female)
#    return Population(females,males,capacity=len(females+males))

class WeightedRandomGeneratorWithReplacement(object):
    def __init__(self, weights):
        self.totals = []
        running_total = 0

        for w in weights:
            running_total += w
            self.totals.append(running_total)

    def next(self):
        rnd = npr.random() * self.totals[-1]
        return bisect.bisect_right(self.totals, rnd)

    def __call__(self):
        return self.next()

class WeightedRandomGeneratorWithoutReplacement(WeightedRandomGeneratorWithReplacement):
    def __init__(self, weights):
        WeightedRandomGeneratorWithReplacement.__init__(self, weights)
        self.selected = []

    def next(self):
        if len(self.totals) == len(self.selected):
            raise ValueError, "generator is empty"
        rnd = npr.random() * self.totals[-1]
        i = bisect.bisect_right(self.totals, rnd)
        if i not in self.selected:
            self.selected.append(i)
            return i
        else:
            return self.next()

class Malekiller(object):
    def __init__(self, lethality=None, transmission_rate=None, parameters=_PARAMETERS):
        if lethality==None:
            self.lethality = parameters['mk_lethality']
        else:
            self.lethality = lethality
        if transmission_rate==None:
            self.transmission_rate = parameters['mk_transmission_rate']
        else:
            self.transmission_rate = transmission_rate
           
    def __str__(self):
        s = """\
            male killer
            -----------
               |-- lethality: {self.lethality}
               |-- transmission rate: {self.transmission_rate}
            """.format(self=self)
        return textwrap.dedent(s)
    
    def __repr__(self):
        return "Malekiller( lethality={self.lethality}, transmission_rate={self.transmission_rate} )".format(self=self)

class Individual(object):
    def __init__(self, genotype, infection=None, fitness=None, parameters=_PARAMETERS):
        self.genotype = genotype
        self.infection = self.mk = infection
        self.infected = self.infection is not None             # bool value so you can test for infection with 'if self.infected:...'
        if fitness == None:
            fitness = parameters['base_fitness']
        self.fitness = fitness
        self._alleles = parameters['alleles'] 
        self._phenomap = parameters['phenotype_map']
        self.phenotype = self.get_phenotypic_value()
        self._mutation_rate = parameters['mutation_rate']
        self._parameters = parameters
        
    def __str__(self):
        s = """\
            individual
            ----------
                |-- {infection}
                |-- phenotype: {self.phenotype}
                |-- fitness: {self.fitness}
                |-- genotype: {gtype}
            """.format( self=self, infection=self.infection2string(), gtype=self.genotype2string() )
        return textwrap.dedent(s)
        
    def __repr__(self):
        return "Individual( genotype={s.genotype}, infection={s.infection!r}, fitness={s.fitness}, parameters={s._parameters!r})".format(s=self)        
    
    def infection2string(self):
        if self.infected:
            return "infected with {}".format( treeindent(str(self.infection),22) )
        else:
            return "uninfected"
    
    def genotype2string(self):
        return '{g[0]} || {g[1]}'.format(g=self.genotype)
        
    def get_phenotypic_value(self):
        """Take the average."""
        a1,a2 = self.genotype
        return ( self._phenomap[a1] + self._phenomap[a2] ) / 2.
    
    def update_fitness(self, factor):
        self.fitness *= factor
    
    def gamete(self):
        mutation = 0
        g = random_choice( self.genotype )
        # mutation:
        if npr.random() < self._mutation_rate:
            p_jump = self.parameters['probability_of_large_effect_mutation']
            g = mutate(g, self._alleles, mode='mixed', p_jump=p_jump)
            mutation = 1
        return g, mutation

class Female(Individual):
    def __init__(self, genotype, infection=None, fitness=None, parameters=_PARAMETERS):
        Individual.__init__(self, genotype=genotype, infection=infection, fitness=fitness, parameters=parameters)
        self.sex = 'female'
        self.sexratio = math.floor(self.phenotype)
        self.fecundity = parameters['female_fecundity']
        self._realloceff = parameters['reallocation_efficiency']
    
    def __repr__(self):
        return "Female( genotype={s.genotype}, infection={s.infection!r}, fitness={s.fitness}, parameters={s._parameters!r})".format(s=self)
    
    def __str__(self):
        s = """\
            Female
            ------
                |-- {infection}
                |-- phenotype: {self.phenotype}
                |-- sex ratio: {self.sexratio}
                |-- fitness: {self.fitness}
                |-- genotype: {gtype}
            """.format( self=self, infection=self.infection2string(), gtype=self.genotype2string() )
        return textwrap.dedent(s)
    
    def produce_brood_with(self, male):
        mut_count = 0
        sons = []
        daughters = []
        pars = self._parameters
        for i in range(1,self.fecundity+1):
            # offspring genotype:
            egg, m1   = self.gamete()
            sperm, m2 = male.gamete()
            mut_count += m1+m2
            # offspring infection:
            infection = None
            if self.infected:
                if npr.random() < self.mk.transmission_rate:               # random transmission of infection
                    infection = self.mk
            # offspring sex:
            if i <= self.sexratio:
                sons.append( Male([egg,sperm], infection, parameters=pars) )
            else:
                daughters.append( Female([egg,sperm], infection, parameters=pars) )
        if self.infected:
            sons,daughters = self.male_killing(sons,daughters)
        male.number_of_matings += 1
        return sons, daughters, mut_count
    
    def male_killing(self, sons, daughters):
        n = len(sons)
        sons[:] = (s for s in sons if s.infected and not npr.random()<self.mk.lethality)
        body_count = n - len(sons)
        for offspring in sons+daughters:
            fc = fitness_compensation(body_count, self.fecundity, self._realloceff, rmin=None)
            offspring.update_fitness( fc )
        return sons,daughters
    
class Male(Individual):
    def __init__(self, genotype, infection=None, fitness=None, parameters=_PARAMETERS):
        Individual.__init__(self, genotype=genotype, infection=infection, fitness=fitness, parameters=parameters)
        self.sex = 'male'
        self.number_of_matings = 0
        
    def __str__(self):
        s = """\
            Male
            -----
                |-- {infection}
                |-- number of matings: {self.number_of_matings}
                |-- fitness: {self.fitness}
                |-- genotype: {gtype}
            """.format( self=self, infection=self.infection2string(), gtype=self.genotype2string() )
        return textwrap.dedent(s)
        
    def __repr__(self):
        return "Male( genotype={s.genotype}, infection={s.infection!r}, fitness={s.fitness}, parameters={s._parameters!r} )".format(s=self)

class Population(object):                             
    def __init__(self, males, females, generation=0, record=None, parameters=_PARAMETERS):
        self.males = males
        self.females = females
        self.size = len(males+females)        # distinguish between size and capacity
        self.generation = generation
        self.capacity = parameters['population_capacity']
        self._alleles = parameters['alleles']
        self._female_fecundity = parameters['female_fecundity']
        self._reallocation_efficiency = parameters['reallocation_efficiency']
        self._parameters = parameters
        
        if record == None:
            record_keys = ['generation','n','n_males','n_females','n_mk','n_mk_males','n_mk_females','n_mutations','average_sex_ratio','allele_count','genotype_count']
            self.record = {}
            for k in record_keys:
                self.record[k] = []
            self.update_record(mut_count=0)
        else:
            self.record = record
            
    def __str__(self):
        rec = {}                          # extract current state dictionary from record for easy access
        for k,l in self.record.items():
            rec[k] = l[-1]
        s = """\
            Population
            ----------
                |-- generation: {rec[generation]}
                |-- size: {rec[n]}  (capacity {s.capacity})
                |-- number of males: {rec[n_males]}  ({rec[n_mk_males]} infected)
                |-- number of females: {rec[n_females]}  ({rec[n_mk_females]} infected)
                |-- total infected: {rec[n_mk]}
                |-- average sex ratio: {rec[average_sex_ratio]}
            """.format( rec=rec, s=self )
        return textwrap.dedent(s)
    
    def allele_count(self, gcount=None):
        if gcount == None:
            gcount = self.genotype_count()
        return np.sum(gcount,0) + np.sum(gcount,1)
    
    def genotype_count(self):
        alleles = self._alleles
        N = len(alleles)
        gcount = np.zeros((N,N), dtype=int)
        for i in self.males+self.females:
            a1,a2 = i.genotype
            gcount[alleles.index(a1),alleles.index(a2)] += 1
        return gcount
    
    def average_sex_ratio(self):
        """Calculate average sex ratio (asr) among females."""
        value = 0.
        for f in self.females:
            value += f.sexratio
        return value/len(self.females)
    
    def mk_count(self):
        mk_males = mk_females = 0
        for m in self.males:
            if m.infected:
                mk_males += 1
        for f in self.females:
            if f.infected:
                mk_females += 1
        return mk_males, mk_females
    
    def tally(self):
        """Calculate average sex ratio, mk numbers, and genotype frequencies; but avoid iterating over population multiple times."""
        alleles = self._alleles
        nA = len(alleles)
        asr = 0.                                # average sex ratio among females
        mk_males = mk_females = 0               # numbers of infected males and females
        acount = [0 for a in alleles]           # allele count
        gcount = np.zeros((nA,nA), dtype=int)   # genotype count
        for m in self.males:
            if m.infected:
                mk_males += 1
            a1,a2 = m.genotype
            gcount[ alleles.index(a1), alleles.index(a2) ] += 1
        for f in self.females:
            asr += f.sexratio
            if f.infected:
                mk_females += 1
            a1,a2 = f.genotype
            gcount[ alleles.index(a1), alleles.index(a2) ] += 1
        asr /= len(self.females)
        acount = self.allele_count(gcount)
        return asr, mk_males, mk_females, acount, gcount            
    
    def update_record(self, mut_count):
        try:
            last_stored = self.record['generation'][-1]
        except IndexError:                     # this Error is raised when the record is still empty
            last_stored = -1
        if self.generation != last_stored:     # only update record if the generation state has not been stored before
            self.record['generation'].append( self.generation )
            self.record['n'].append( self.size )
            self.record['n_males'].append( len(self.males) )
            self.record['n_females'].append( len(self.females) )
            self.record['n_mutations'].append( mut_count )
            sr, mkM, mkF, acount, gcount = self.tally()
            self.record['average_sex_ratio'].append( sr )
            self.record['n_mk_males'].append( mkM )
            self.record['n_mk_females'].append( mkF )
            self.record['n_mk'].append( mkM+mkF )
            self.record['allele_count'].append( acount )
            self.record['genotype_count'].append( gcount )
        
    def is_polymorphic(self):
        acount = self.record['allele_count'][-1]         # most current allele count
        maxcount = 2*self.capacity
        return maxcount not in acount
        
    def produce_offspring(self):
        mut_count = 0
        if self.males == []:
            raise ValueError, 'population extinct due to lack of males'
        if self.females == []:
            raise ValueError, 'population extinct due to lack of females'
            
        males = self.males[:]
        npr.shuffle(males)
        females = self.females[:]
        npr.shuffle(females)
        temp_male_pool = []            # after having mated, males go into this pool
        offspring = []                 # all offspring is collected here
        
        for f in females:
            try:
                m = males.pop()
            except IndexError:                  # if there are no unmated males left, 
                males = temp_male_pool[:]       # put all males from the temporary pool back
                temp_male_pool = []             # and empty the temporary pool again
                npr.shuffle(males)
                m = males.pop()
            sons,daughters,mut = f.produce_brood_with(m)
            mut_count += mut
            temp_male_pool.append(m)
            offspring.extend(sons+daughters)
        return offspring, mut_count
        
    def selection(self, offspring):
#        print "offspring size:", len(offspring)
        males = []
        females = []
        weights = []
        for o in offspring:
            weights.append(o.fitness)
        wrg = weighted_random_generator(weights, replacement=False)
        i = 0
        while i < self.capacity:
            try:
                individual = offspring[ wrg() ]
                if individual.sex == 'female':
                    females.append(individual)
                else:
                    males.append(individual)
                i += 1
            except IndexError:           # all offspring have made it into the next generation already (offspring number < population capacity)
                break
        return males, females
    
    def next(self):
        offspring, mut_count = self.produce_offspring()
        self.males, self.females = self.selection(offspring)
        self.size = len(self.males+self.females)
        self.generation += 1
        self.update_record(mut_count)
    
    def run(self, n, force=False, timing=False):
        """If force is False then n is maximal generation count but iteration stops if the population
        is no longer polymorphic."""
        if timing:
            g0 = self.generation
            t0 = time.time()
        if force:
            for i in range(n):
                self.next()
        else:
            g = 0
            while self.is_polymorphic() and g<n:
                self.next()
                g += 1
        if timing:
            print "Iterated {} generations in {:.2f} seconds".format(self.generation-g0, time.time()-t0)
            
    def visualize(self, plots=['asr','nM','mk','mut','acount','inds']):
        F = self._female_fecundity
        labels = dict(asr=dict(rec='average_sex_ratio',xlabel='generation',ylabel='average sex ratio',style='k-',ylims=[0,F]), \
                      nM=dict(rec='n_males',xlabel='generation',ylabel='number of males',style='g-',ylims=[0,self.capacity]), \
                      mk=dict(rec='n_mk',xlabel='generation',ylabel='MK infected',style='r-',ylims=[0,self.capacity]), \
                      mut=dict(rec='n_mutations',xlabel='generation',ylabel='number of mutations',style='m-'), \
                      acount=dict(rec='allele_count',xlabel='generation',ylabel='alleles'), \
                      inds=dict(rec='individuals',xlabel='allele 1',ylabel='allele 2') )
        if not isinstance(plots, list):
            plots = list(plots)
        if 'inds' in plots:
            plots.remove('inds')
        figs = []
        g = np.array(self.record['generation'])
        for p in plots:
            f = plt.figure()
            ax = f.add_subplot(111)
            figs.append(f)
            rec = labels[p]['rec']
            xl = labels[p]['xlabel']
            yl = labels[p]['ylabel']
            
            a = np.array(self.record[rec])
            if p == 'mut':
                a = np.cumsum(a)
                labels[p]['ylims']= [0,np.max(a)]
            if p == 'acount':
                a = np.transpose( np.array(a,dtype=float) )
                a /= (2*self.capacity)                         # normalize (so that values range between 0 and 1)
                ma = np.ma.array(a,mask=(a==0))            # mask zero values
                cmap = mpl.cm.jet                          # use jet color scheme
                cmap.set_bad('w')                          # use white for zero values
                im = ax.imshow(ma,cmap=cmap,vmin=0.,vmax=1.,origin='lower',interpolation='nearest',aspect='auto')
                cb = plt.colorbar(im)
                cb.set_label('allele frequency')
            elif p == 'inds':
                pass
            else:
                ylims = labels[p]['ylims']
                s = labels[p]['style']
                ax.plot(g,a,s)
                ax.set_ylim(*ylims)
            ax.set_xlabel(xl)
            ax.set_ylabel(yl)
        return plots,figs
        
    def genotype_show(self, g=-1):
        f = plt.figure()
        ax = f.add_subplot(111)
        a = np.array( self.record['genotype_count'][g], dtype=float )
        a /= self.capacity                                      # normalize
        ma = np.ma.array(a, mask=(a==0))
        cmap = mpl.cm.jet
        cmap.set_bad('w')
        im = ax.imshow(ma,cmap=cmap,vmin=0.,vmax=1.,origin='lower',interpolation='nearest',aspect='auto')
        cb = plt.colorbar(im)
        cb.set_label('frequency')
        ax.set_xlabel('chromosome 1')
        ax.set_ylabel('chromosome 2')
    
    def store_current_state(self, filename='population.dat'):
        pass
        
    
if __name__ == '__main__':
    npr.seed(21112016)
    
    pars = init_parameters(female_fecundity=10, mk_lethality=1., \
                           base_fitness=1., reallocation_efficiency=1., \
                           population_capacity=1000)
    print_parameters(pars)
    
    mk = Malekiller(parameters=pars)                      # use provided parameters (recommended usage)
#    mk2 = Malekiller()                                    # use module default parameters (good for debugging)
#    mk3 = Malekiller(lethality=0.5)                       # use default parameters but override lethality
#    mk4 = Malekiller(transmission_rate=0.25, parameters=pars)     # use provided parameters but override lethality
#    print mk, mk2, mk3, mk4
#    
#    f = Female(['a4','a7'], mk)
#    m = Male(['a5','a6'], mk)
#    print f
#    print m
#    sons,daughters,mut = f.produce_brood_with(m)
    
    pop = random_population(1000)
    print pop
    
    pop.run(100)
    print pop
    pop.visualize()
    plt.show()
#    pop = uniform_population()
##    females = [Female(['a1','a2'],'uninfected') for i in range(30)]
##    males = [Male(['a1','a2'],'uninfected') for i in range(30)]
##    pop = Population(females,males)
#    pop.run(10, force=True, timing=True)
#    
#    plots,figs = pop.visualize()    
#    plt.show()
